package SPH.RadixSort;

//import static com.jogamp.opencl.CLMemory.Mem.READ_WRITE;

import java.io.IOException;
//import java.nio.ByteBuffer;
//import java.nio.IntBuffer /*FloatBuffer*/;
import java.nio.IntBuffer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

import SPH.SPHSolver;

import com.jogamp.opencl.CLBuffer;
import com.jogamp.opencl.CLCommandQueue;
import com.jogamp.opencl.CLContext;
import com.jogamp.opencl.CLDevice;
import com.jogamp.opencl.CLEvent;
import com.jogamp.opencl.CLEventList;
import com.jogamp.opencl.CLKernel;
import com.jogamp.opencl.CLMemory;
import com.jogamp.opencl.CLProgram;


public class RadixSortCL {
	 int _sortSize;
     int _numKeyBits;
     int _numIssueGroups;
     int _scanBlockSize;
     int _numScanGroups;

     public static final int GROUP_SIZE = 128;
     public static final int ELEMENTS_PER_WORK_ITEM = 4;
     public static final int SCAN_WORKGROUP_SIZE = 128;
     public static final int BLOCK_SUM_WORKGROUP_SIZE = 128;
     public static final int BITS_PER_PASS = 4;
     final int _groupSize;
     final int _elementsPerWorkItem;
     /* Set to true if the intermediate buffer has been initialized */
     boolean _intermediateBufferInitialized;
     boolean _environmentInitialised;
     List<String> _kernelNames;
     HashMap<String, CLKernel> kernels;
     CLBuffer<IntBuffer /*FloatBuffer*/> _scratchBuffer; 
     // Local histogram buffer that stays local
     CLBuffer<IntBuffer> localHistogramsBuffer;
     // Radix major layout local hist buffer that will be scanned, before and after
     // (these can probably be merged to do a scan-in-place
     CLBuffer<IntBuffer> localHistogramsBufferGlobalSummed;
     CLBuffer<IntBuffer> localHistogramsBufferGlobalSummedPostScan;
     CLBuffer<IntBuffer> scanIntermediateBlockSumBuffer;

     CLEvent _event;
     CLContext context;
     CLCommandQueue queue;
     CLProgram program;
     //binaries;

     int status;
     boolean _profile;
     
     /*
      * \brief OpenCL related initialisations are done here.
      *        Context, Device list, Command Queue are set up.
      *        Calls are made to set up OpenCL memory buffers that this program uses
      *        and to load the programs into memory and get kernel handles.
      */
     void initializeAcceleratorEnvironment() throws IOException{
    	// Create and build CL program and load all kernels
         // an array with available devices
      	CLDevice[] devices = context.getDevices();
      	if( devices.length == 0 ) {
      		System.out.println( "No CL devices have been found" );
      		System.exit( 1 );
      	}
      	
      	program = context.createProgram(SPHSolver.class.getResourceAsStream("/resource/radixSortkernels.cl")).build();
     }
     /*
      * \brief Set up OpenCL memory buffers used by this program
      */
     void initializeAcceleratorBuffers(int elementCount, int numBlocks, int numScanGroups){
    	 localHistogramsBuffer = context.createIntBuffer((1<<BITS_PER_PASS)*numBlocks, CLMemory.Mem.READ_ONLY);
    	 localHistogramsBufferGlobalSummed = context.createIntBuffer((1<<BITS_PER_PASS)*numBlocks, CLMemory.Mem.READ_ONLY);
    	 localHistogramsBufferGlobalSummedPostScan = context.createIntBuffer((1<<BITS_PER_PASS)*numBlocks, CLMemory.Mem.READ_ONLY);
    	 scanIntermediateBlockSumBuffer = context.createIntBuffer(numScanGroups, CLMemory.Mem.READ_ONLY);
     }
     /*
      * \brief Load and build OpenCL program and get kernel handles
      */
     void initializeAcceleratorKernels(){
    	 for(int i = 0; i < _kernelNames.size(); i++)
	         kernels.put(_kernelNames.get(i),
	        		 program.createCLKernel(_kernelNames.get(i)));
	 }
     /*
      * \brief Set up intermediate OpenCL buffer that will be used for computation.
      */
     void initializeAcceleratorDataBuffer(int elementCount, int numBlocks, int numScanGroups){
    	 _scratchBuffer = context.createIntBuffer /*FloatBuffer*/(elementCount,  CLMemory.Mem.READ_ONLY);
    	 _intermediateBufferInitialized = true;
     }
     /**
      * Performs a prefix sum of inputScanArray into outputScanArray using scanBlockSums as an intermediate
      * buffer for the sums of individual work groups during the process.
      *
      * Current version uses the CPU to do block sums as it's a simple operation but this has excessive overhead, really.
      * The propagation of block sums is done on the GPU, but really everything should be.
      * For very large arrays the block sum computation will need a tree of blocks rather than a single block with a two
      * step global process, so that should be done carefully eventually.
      * End event is the class variable _event.
      */
     void scanLargeArrays(CLKernel scanLargeArraysKernel,
             CLKernel sumBlockSumsKernel,
             CLKernel scanPropagateBlockSumsKernel,
             CLEventList sortStartEvents)
     {
    	 queue.put1DRangeKernel(scanLargeArraysKernel, 0, _numScanGroups*SCAN_WORKGROUP_SIZE, SCAN_WORKGROUP_SIZE, sortStartEvents);
    	 if(_profile){
    		 try {
				_event.wait();
			} catch (InterruptedException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}
    		long submit = _event.getProfilingInfo(CLEvent.ProfilingCommand.SUBMIT);
     		long start =  _event.getProfilingInfo(CLEvent.ProfilingCommand.START);
     		long end =  _event.getProfilingInfo(CLEvent.ProfilingCommand.END);
     		System.out.println("Elapsed time for local scan kernel execution: " + (double)(end - start)/1e6 + "ms and from submit to end: " + (double)(end - submit)/1e6 );
    	 }
 		// Reprepare sort event structure
        sortStartEvents = new CLEventList( _event );
        
        queue.put1DRangeKernel(sumBlockSumsKernel, 0, BLOCK_SUM_WORKGROUP_SIZE, BLOCK_SUM_WORKGROUP_SIZE, sortStartEvents);
        if(_profile){
        	try {
				_event.wait();
			} catch (InterruptedException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}
        	long submit = _event.getProfilingInfo(CLEvent.ProfilingCommand.SUBMIT);
     		long start =  _event.getProfilingInfo(CLEvent.ProfilingCommand.START);
     		long end =  _event.getProfilingInfo(CLEvent.ProfilingCommand.END);
     		System.out.println("Elapsed time for local scan kernel execution: " + (double)(end - start)/1e6 + "ms and from submit to end: " + (double)(end - submit)/1e6 );
        }
        sortStartEvents = new CLEventList( _event );
        if( _numScanGroups > 1 ){
        	queue.put1DRangeKernel(scanPropagateBlockSumsKernel, 0, SCAN_WORKGROUP_SIZE*(_numScanGroups - 1), SCAN_WORKGROUP_SIZE, sortStartEvents);
        	if(_profile){
            	try {
    				_event.wait();
    			} catch (InterruptedException e) {
    				// TODO Auto-generated catch block
    				e.printStackTrace();
    			}
            	long submit = _event.getProfilingInfo(CLEvent.ProfilingCommand.SUBMIT);
         		long start =  _event.getProfilingInfo(CLEvent.ProfilingCommand.START);
         		long end =  _event.getProfilingInfo(CLEvent.ProfilingCommand.END);
         		System.out.println("Elapsed time for local scan kernel execution: " + (double)(end - start)/1e6 + "ms and from submit to end: " + (double)(end - submit)/1e6 );
            }
            sortStartEvents = new CLEventList( _event );
        }
     }
     /*
      * \brief The internal sort function.
      * Data assumed to be in the correct start buffer to end up in outputBuffer at the end.
      */
     void sortInternal(
             CLBuffer<IntBuffer /*FloatBuffer*/> scratchBuffer,
             CLBuffer<IntBuffer /*FloatBuffer*/> outputBuffer,
             boolean oddPassCount,
             CLEventList sortStartEvents ){
    	 CLBuffer<IntBuffer /*FloatBuffer*/> buffer1;
    	 CLBuffer<IntBuffer /*FloatBuffer*/> buffer2;
    	 if( oddPassCount )
         {
             buffer1 = scratchBuffer;
             buffer2 = outputBuffer;
         } else {
             buffer1 = outputBuffer;
             buffer2 = scratchBuffer;
         }
    	 int elementsPerBlock = _groupSize*_elementsPerWorkItem;
    	// Local memory storage
    	 int totalLocalMemoryWords = elementsPerBlock + 16;
    	// Set static kernel arguments out of main loop
    	 CLKernel radixSortLocalKernel = kernels.get("RadixSortLocal");
         radixSortLocalKernel.setArg(1, localHistogramsBufferGlobalSummed);
         radixSortLocalKernel.setArg(2, localHistogramsBuffer);
         radixSortLocalKernel.setArg(4, _numIssueGroups);
         // *2 for initial counters and then for prefix sums
         // maybe that can compress into the same array
         radixSortLocalKernel.setArg(5,totalLocalMemoryWords);

         CLKernel globalSortKernel = kernels.get("RadixSortGlobal");
         globalSortKernel.setArg(1, localHistogramsBufferGlobalSummedPostScan);
         globalSortKernel.setArg(2, localHistogramsBuffer);
         globalSortKernel.setArg(5, (3*(1<<BITS_PER_PASS)));


         CLKernel scanLargeArraysKernel = kernels.get("ScanLargeArrays");
         scanLargeArraysKernel.setArg(0, localHistogramsBufferGlobalSummedPostScan);
         scanLargeArraysKernel.setArg(1, localHistogramsBufferGlobalSummed);
         scanLargeArraysKernel.setArg(2, _scanBlockSize);
         scanLargeArraysKernel.setArg(3, _scanBlockSize);
         scanLargeArraysKernel.setArg(4, _numIssueGroups*(1<<BITS_PER_PASS));
         scanLargeArraysKernel.setArg(5, scanIntermediateBlockSumBuffer);

         CLKernel sumBlockSumsKernel = kernels.get("SumBlockSums");
         sumBlockSumsKernel.setArg(0, scanIntermediateBlockSumBuffer);
         sumBlockSumsKernel.setArg(1, _numScanGroups);
         // TODO: This number's a bit random. Need to make it based off something
         sumBlockSumsKernel.setArg(2, 2048);

         CLKernel scanPropagateBlockSumsKernel = kernels.get("ScanPropagateBlockSums");
         scanPropagateBlockSumsKernel.setArg(0, localHistogramsBufferGlobalSummedPostScan);
         scanPropagateBlockSumsKernel.setArg(1, _scanBlockSize);
         scanPropagateBlockSumsKernel.setArg(2, _numIssueGroups*(1<<BITS_PER_PASS));
         scanPropagateBlockSumsKernel.setArg(3, scanIntermediateBlockSumBuffer);
         // Main loop over sets of bits
         for( int bitIndex = 0; bitIndex < _numKeyBits; bitIndex += BITS_PER_PASS  ){
             if ( _profile )
            	 System.out.println("Bit index: " + bitIndex + "\n");
             // Stage 1 and 2 from Satish paper:
             // Sort blocks of array in local memory of separate work groups
             // Output the sorted sublists and histograms representing the
             // number of each value in the sublist
             radixSortLocalKernel.setArg(0, buffer1);
             radixSortLocalKernel.setArg(3, bitIndex);
             // TODO enqueueNDRangeKernel( radixSortLocalKernel, cl::NullRange, global, local, &sortStartEvents, &_event); Used in c++ example
             queue.put1DRangeKernel(radixSortLocalKernel, 0, _numIssueGroups*_groupSize, _groupSize);
          
             // Reprepare sort event structure
             sortStartEvents = new CLEventList( _event );
            
             if(_profile){
             	try {
     				_event.wait();
     			} catch (InterruptedException e) {
     				// TODO Auto-generated catch block
     				e.printStackTrace();
     			}
             	long submit = _event.getProfilingInfo(CLEvent.ProfilingCommand.SUBMIT);
          		long start =  _event.getProfilingInfo(CLEvent.ProfilingCommand.START);
          		long end =  _event.getProfilingInfo(CLEvent.ProfilingCommand.END);
          		System.out.println("Elapsed time for local scan kernel execution: " + (double)(end - start)/1e6 + "ms and from submit to end: " + (double)(end - submit)/1e6 );
             }
             // State 3 from Satish paper:
             // Perform prefix sum over histogram bins
             scanLargeArrays(
                     scanLargeArraysKernel,
                     sumBlockSumsKernel,
                     scanPropagateBlockSumsKernel,
                     sortStartEvents
                     );

             // Stage 4 from Satish paper:
             // Do radix shuffle based on the prefix summed offsets
             globalSortKernel.setArg(0, buffer1);
             globalSortKernel.setArg(3, buffer2);
             globalSortKernel.setArg(4, bitIndex);
             // 2x because we're storing two histograms - local and global
             queue.put1DRangeKernel(globalSortKernel, 0, _numIssueGroups*_groupSize, _groupSize, sortStartEvents);
             if(_profile){
              	try {
      				_event.wait();
      			} catch (InterruptedException e) {
      				// TODO Auto-generated catch block
      				e.printStackTrace();
      			}
              	long submit = _event.getProfilingInfo(CLEvent.ProfilingCommand.SUBMIT);
           		long start =  _event.getProfilingInfo(CLEvent.ProfilingCommand.START);
           		long end =  _event.getProfilingInfo(CLEvent.ProfilingCommand.END);
           		System.out.println("Elapsed time for local scan kernel execution: " + (double)(end - start)/1e6 + "ms and from submit to end: " + (double)(end - submit)/1e6 );
             }
             // Swap buffers
             CLBuffer<IntBuffer /*FloatBuffer*/> tempBuffer = buffer1;
             buffer1 = buffer2;
             buffer2 = tempBuffer;

             // Reprepare sort event structure for next iteration
             sortStartEvents.release();
             sortStartEvents =  _event ;
         }
     }
     
     public RadixSortCL(){
    	 _sortSize = 0;
         _numKeyBits = 32;
         _numIssueGroups = 0;
         _scanBlockSize = SCAN_WORKGROUP_SIZE*2;
         _numScanGroups = 0;
         _groupSize = GROUP_SIZE;
         _elementsPerWorkItem = ELEMENTS_PER_WORK_ITEM;
         _intermediateBufferInitialized = false;
         _environmentInitialised = false;
         _kernelNames = new ArrayList<String>();
         _kernelNames.add("RadixSortLocal");
         _kernelNames.add("ScanLargeArrays");
         _kernelNames.add("ScanPropagateBlockSums");
         _kernelNames.add("RadixSortGlobal");
         _kernelNames.add("SumBlockSums");
         _profile = false;
     }
     
     public void initializeSort(
             final CLContext contextL,
             final CLCommandQueue queue,
             int sortSize ) {
    
    	 initializeSort( contextL, queue, sortSize, 32, true ); 
     }
     
     public void initializeSort(
             final CLContext contextL,
             final CLCommandQueue queue,
             int sortSize,
             int keySize,
        	 boolean initializeInternalBuffer )
     {
    	 _sortSize = sortSize;
         _numKeyBits = keySize;
         _numIssueGroups = ( (_sortSize / ELEMENTS_PER_WORK_ITEM)/GROUP_SIZE );
         _numScanGroups = (int)((float)(_numIssueGroups)*(1<<BITS_PER_PASS)/_scanBlockSize);
         
         if( !_environmentInitialised )
         {
             context = contextL;
             this.queue   = queue;
             try {
				initializeAcceleratorEnvironment();
			} 
            catch (IOException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			}
             initializeAcceleratorKernels();
             _environmentInitialised = true;
         }
         initializeAcceleratorBuffers( _sortSize, _numIssueGroups, _numScanGroups );
         // Create the internal data buffer if requested
         
         if( initializeInternalBuffer )
             initializeAcceleratorDataBuffer( _sortSize, _numIssueGroups, _numScanGroups );
     }

     /**
      * Sort inputBuffer into outputBuffer.
      * If mutableInputBuffer set then the inputBuffer is allowed to be used as scratch.
      * If not an internal scratch buffer will be used (and created if not already).
      */
     void sort( CLBuffer<IntBuffer /*FloatBuffer*/> inputBuffer, CLBuffer<IntBuffer /*FloatBuffer*/> outputBuffer, boolean mutableInputBuffer, final CLEventList events )
     {
    	 CLEvent returnEvent = null;
         sort( inputBuffer, outputBuffer, mutableInputBuffer, returnEvent, events );
     }
     /**
      * Sort inputBuffer into outputBuffer.
      * If mutableInputBuffer set then the inputBuffer is allowed to be used as scratch.
      * If not an internal scratch buffer will be used (and created if not already).
      */
     void sort( CLBuffer<IntBuffer /*FloatBuffer*/> inputBuffer, CLBuffer<IntBuffer /*FloatBuffer*/> outputBuffer, boolean mutableInputBuffer, CLEvent endEvent, final CLEventList events)
     {
    	 int numPasses = _numKeyBits/BITS_PER_PASS;
         boolean oddPassCount = ((numPasses & 1) == 1) ? true:false;
         boolean requireScratchBuffer = (numPasses>1) && !mutableInputBuffer;
         
         CLEventList sortStartEvents;
         sortStartEvents = events;

         // If intermediate scratch buffer hasn't already been initialized, create it
         if( requireScratchBuffer && !_intermediateBufferInitialized )
             initializeAcceleratorDataBuffer( _sortSize, _numIssueGroups, _numScanGroups );
         
         if( requireScratchBuffer )
         {
             // If we need a scratch buffer then copy out of the input buffer to start
             // and use the scratch buffer in the sort

             // If we have an even number of passes start in the output buffer, otherwise start in the scratch buffer
             // Start with copy to initialize this
             if( oddPassCount ) {
                 // Enqueue a copy and make its event the one necessary for the sort
                 queue.putCopyBuffer( inputBuffer, _scratchBuffer, 0, 0, _sortSize, sortStartEvents, new CLEventList( _event ) );
                 sortStartEvents = new CLEventList( _event );
             } 
             else {
                 // Enqueue a copy and make its event the one necessary for the sort
                 queue.putCopyBuffer( inputBuffer, outputBuffer, 0, 0, _sortSize, sortStartEvents, new CLEventList( _event ) );
                 sortStartEvents = new CLEventList( _event );
             }
             sortInternal( _scratchBuffer, outputBuffer, oddPassCount, sortStartEvents );
         } 
         else {
        	 // sortInternal( _scratchBuffer, outputBuffer, oddPassCount, sortStartEvents ); No such line in C++ example
        	 if( !oddPassCount ){
                 // Enqueue a copy and make its event the one necessary for the sort
                 queue.putCopyBuffer( inputBuffer, outputBuffer, 0, 0, _sortSize, sortStartEvents, new CLEventList( _event ) );
                 // Reprepare sort event structure
                 sortStartEvents = new CLEventList( _event );
             }
        	 sortInternal( inputBuffer, outputBuffer, oddPassCount, sortStartEvents );
         }
         // Return the termination event to the caller
         endEvent = _event;
     }
     
     public void sort( CLBuffer<IntBuffer /*FloatBuffer*/> IOBuffer )
     {
    	 sort( IOBuffer, new CLEventList() );
     }
     
     void sort( CLBuffer<IntBuffer /*FloatBuffer*/> IOBuffer, final CLEventList events )
     {
         int numPasses = _numKeyBits/BITS_PER_PASS;
         boolean oddPassCount = ((numPasses & 1) == 1)? true:false;

         // Always need a scratch buffer with only one input buffer
         if( !_intermediateBufferInitialized )
             initializeAcceleratorDataBuffer( _sortSize, _numIssueGroups, _numScanGroups );
         CLEventList sortStartEvents = events;
         // If we have an odd number of passes, then we must copy the data into the scratch buffer to start
         if( oddPassCount )
         {
             // Enqueue a copy and make its event the one necessary for the sort
             queue.putCopyBuffer( IOBuffer, _scratchBuffer, 0, 0, _sortSize, sortStartEvents, new CLEventList( _event ) );
            
             // Reprepare sort event structure
             sortStartEvents = new CLEventList( _event );
         }

         sortInternal( _scratchBuffer, IOBuffer, oddPassCount, sortStartEvents );
     }
     
     void sort( CLBuffer<IntBuffer /*FloatBuffer*/> IOBuffer, CLEvent endEvent , final CLEventList events ){
    	 sort( IOBuffer, events );

         // Return the termination event to the caller
         endEvent = _event;
     }
     /**
      * Wait for completion of last operation.
      * Blocks until the sort is completed.
      */
     public void WaitSort(){
         try {
			_event.wait();
		} catch ( InterruptedException e ) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
     } // wait
}
